import tensorflow as tf
from tensorflow.keras import Model
from tensorflow.keras.layers import Dense, Flatten, Conv2D, Input, MaxPooling2D, Dropout, concatenate
from tensorflow.keras.optimizers import Nadam

class Net(Model):
    def __init(self, input_shape, output_num):
        super(Net, self).__init__()
        self.input_shape = input_shape
        self.adam = None
        self.model = None
        self.create_model(output_num)

    def create_model(self, output_num):
        activation = 'relu'

        img_input = Input(shape=self.input_shape[0])

        img_stack = Conv2D(16, 3, name='conv0', padding='same', activation=activation)(img_input)
        img_stack = MaxPooling2D(pool_size=2)(img_stack)
        img_stack = Conv2D(32, 3, activation=activation, padding='same', name='conv1')(img_stack)
        img_stack = MaxPooling2D(pool_size=2)(img_stack)
        img_stack = Conv2D(32, 3, activation=activation, padding='same', name='conv2')(img_stack)
        img_stack = MaxPooling2D(pool_size=2)(img_stack)
        img_stack = Flatten()(img_stack)
        img_stack = Dropout(rate=0.2)(img_stack)

        # Inject the state input
        state_input = Input(shape=self.input_shape[0])
        merged = concatenate([img_stack, state_input])

        # Add Dense
        merged = Dense(64, activation=activation, name='dense0')(merged)
        merged = Dropout(0.2)(merged)
        merged = Dense(10, activation=activation, name='dense1')(merged)
        merged = Dropout(0.2)(merged)
        merged = Dense(output_num, name='output')(merged)

        adam = Nadam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-08)
        self.model = Model(inputs=[img_input, state_input], outputs=merged)
        self.model.compile(optimizer=adam, loss='mse')
        self.model.summary()
